# K = 3 and d = 4, m = 8, n = 6, C^(A) = 4, C^(B) = 2
import numpy as np
import time


# k = 1, d = 1
A_1_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_1_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 1, d = 2
A_1_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_1_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 1, d = 3
A_1_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_1_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 1, d = 4
A_1_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_1_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 2, d = 1
A_2_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_2_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 2, d = 2
A_2_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_2_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 2, d = 3
A_2_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_2_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 2, d = 4
A_2_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_2_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 3, d = 1
A_3_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_3_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 3, d = 2
A_3_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_3_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 3, d = 3
A_3_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_3_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 3, d = 4
A_3_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_3_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# x
X = np.random.uniform(0.0, 1.0, (8,8,8,8))


def normal_algo_1():
	AB_1_1 = np.kron(A_1_1, B_1_1)
	AB_1_2 = np.kron(A_1_2, B_1_2)
	AB_1_3 = np.kron(A_1_3, B_1_3)
	AB_1_4 = np.kron(A_1_4, B_1_4)
	AB_2_1 = np.kron(A_2_1, B_2_1)
	AB_2_2 = np.kron(A_2_2, B_2_2)
	AB_2_3 = np.kron(A_2_3, B_2_3)
	AB_2_4 = np.kron(A_2_4, B_2_4)
	AB_3_1 = np.kron(A_3_1, B_3_1)
	AB_3_2 = np.kron(A_3_2, B_3_2)
	AB_3_3 = np.kron(A_3_3, B_3_3)
	AB_3_4 = np.kron(A_3_4, B_3_4)

	W_1 = np.concatenate([AB_1_1, AB_2_1, AB_3_1], axis = -1)
	W_2 = np.concatenate([AB_1_2, AB_2_2, AB_3_2], axis = -1)
	W_3 = np.concatenate([AB_1_3, AB_2_3, AB_3_3], axis = -1)
	W_4 = np.concatenate([AB_1_4, AB_2_4, AB_3_4], axis = -1)

	# W_1 to (m_1 * n_1*C)
	W_1 = np.reshape(np.reshape(W_1, (8, 6, -1)), (8, -1))

	# Y is (m_2*m_3*m_4 * n_1*C)
	Y = np.matmul(np.reshape(X, (8, 8*8*8)).T, W_1)

	# Y to (n_1*m_3*m_4 * m_2*C)
	Y = np.reshape(Y, (8, 8*8, 6, -1))
	Y = np.transpose(Y, (2, 1, 0, 3))
	Y = np.reshape(Y, (6*8*8, -1))

	# W_2 to (m_2*C * n_2)
	W_2 = np.reshape(W_2, (8, 6, -1))
	W_2 = np.transpose(W_2, (0, 2, 1))
	W_2 = np.reshape(W_2, (-1, 6))

	# Y is (n_1*m_3*m_4 * n_2)
	Y = np.matmul(Y, W_2)

	# Y to (n_1*n_2*m_4 * m_3)
	Y = np.reshape(Y, (6, 8, 8, 6))
	Y = np.transpose(Y, (0, 3, 2, 1))
	Y = np.reshape(Y, (6*6*8, 8))

	# W_3 to (m_3 * n_3*C)
	W_3 = np.reshape(np.reshape(W_3, (8, 6, -1)), (8, -1))

	# Y is (n_1*n_2*m_4 * n_3*C)
	Y = np.matmul(Y, W_3)

	# Y to (n_1*n_2*n_3 * m_4*C)
	Y = np.reshape(Y, (6*6, 8, 6, -1))
	Y = np.transpose(Y, (0, 2, 1, 3))
	Y = np.reshape(Y, (6*6*6, -1))

	# W_4 to (m_4*C * n_4)
	W_4 = np.reshape(W_4, (8, 6, -1))
	W_4 = np.transpose(W_4, (0, 2, 1))
	W_4 = np.reshape(W_4, (-1, 6))

	# Y is (n_1*n_2*n_3 * n_4)
	Y = np.matmul(Y, W_4)

	return np.reshape(Y, (-1))


def divided_algo_1():
	def single_algo_1(I, W_1, W_2):
		# W_1 to (m_1 * n_1*C)
		W_1 = np.reshape(np.reshape(W_1, (8, 6, -1)), (8, -1))

		# O is (m_2*m_3*m_4 * n_1*C)
		O = np.matmul(np.reshape(I, (8, 8*8*8)).T, W_1)

		# O to (n_1*m_3*m_4 * m_2*C)
		O = np.reshape(O, (8, 8*8, 6, -1))
		O = np.transpose(O, (2, 1, 0, 3))
		O = np.reshape(O, (6*8*8, -1))

		# W_2 to (m_2*C * n_2)
		W_2 = np.reshape(W_2, (8, 6, -1))
		W_2 = np.transpose(W_2, (0, 2, 1))
		W_2 = np.reshape(W_2, (-1, 6))

		# O is (n_1*m_3*m_4 * n_2)
		O = np.matmul(O, W_2)

		# O to (n_1*n_2*m_4, m_3)
		O = np.reshape(O, (6, 8, 8, 6))
		O = np.transpose(O, (0, 3, 2, 1))
		O = np.reshape(O, (6*6*8, 8))

		return O

	def single_algo_1_2(O, W_3, W_4):
		# W_3 to (m_3 * n_3*C)
		W_3 = np.reshape(np.reshape(W_3, (8, 6, -1)), (8, -1))

		# O is (n_1*n_2*m_4, n_3*C)
		O = np.matmul(O, W_3)

		# O to (n_1*n_2*n_3, m_4*C)
		O = np.reshape(O, (6*6, 8, 6, -1))
		O = np.transpose(O, (0, 2, 1, 3))
		O = np.reshape(O, (6*6*6, -1))

		# W_4 to (m_4*C * n_4)
		W_4 = np.reshape(W_4, (8, 6, -1))
		W_4 = np.transpose(W_4, (0, 2, 1))
		W_4 = np.reshape(W_4, (-1, 6))

		# O is (n_1*n_2*n_3, n_4)
		O = np.matmul(O, W_4)

		return np.reshape(O, (-1))

	AB_1_1 = np.kron(A_1_1, B_1_1)
	AB_1_2 = np.kron(A_1_2, B_1_2)
	AB_1_3 = np.kron(A_1_3, B_1_3)
	AB_1_4 = np.kron(A_1_4, B_1_4)
	AB_2_1 = np.kron(A_2_1, B_2_1)
	AB_2_2 = np.kron(A_2_2, B_2_2)
	AB_2_3 = np.kron(A_2_3, B_2_3)
	AB_2_4 = np.kron(A_2_4, B_2_4)
	AB_3_1 = np.kron(A_3_1, B_3_1)
	AB_3_2 = np.kron(A_3_2, B_3_2)
	AB_3_3 = np.kron(A_3_3, B_3_3)
	AB_3_4 = np.kron(A_3_4, B_3_4)

	X_1 = single_algo_1(X, AB_1_1, AB_1_2)
	X_2 = single_algo_1(X, AB_2_1, AB_2_2)
	X_3 = single_algo_1(X, AB_3_1, AB_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_1_2(Y, AB_1_3, AB_1_4)
	X_2 = single_algo_1_2(Y, AB_2_3, AB_2_4)
	X_3 = single_algo_1_2(Y, AB_3_3, AB_3_4)
	Y = X_1 + X_2 + X_3

	return Y


def divided_algo_2():
	def single_algo_2(I, A_1, B_1, A_2, B_2):
		# O is (m_2*m_3*m_4 * C_A)
		O = np.matmul(np.reshape(I, (8, 8*8*8)).T, A_1)

		# O is (m_2*m_3*m_4*n_1 * C_A*C_B)
		O = np.kron(O, B_1)

		# O is (n_1*m_3*m_4*C_B * m_2*C_A)
		O = np.reshape(O, (8, -1, 6, 4, 2))
		O = np.transpose(O, (2, 1, 4, 0, 3))
		O = np.reshape(O, (-1, 8*4))

		# O is (n_1*m_3*m_4 * C_B)
		O = np.matmul(O, np.reshape(A_2, (-1)))
		O = np.reshape(O, (-1, 2))

		# O is (n_1*n_2*m_3*m_4)
		O = np.matmul(O, B_2.T)
		O = np.reshape(O, (6, -1, 6))
		O = np.reshape(np.transpose(O, (0, 2, 1)), (-1))

		return O

	def single_algo_2_2(O, A_3, B_3, A_4, B_4):
		# O to (n_1*n_2*m_4 * m_3)
		O = np.reshape(O, (-1, 8, 8))
		O = np.reshape(np.transpose(O, (0, 2, 1)), (-1, 8))

		# O is (n_1*n_2*m_4 * C_A)
		O = np.matmul(O, A_3)

		# O is (n_1*n_2*m_4*n_3 * C_A*C_B)
		O = np.kron(O, B_3)

		# O is (n_1*n_2*n_3*C_B * m_4*C_A)
		O = np.reshape(O, (-1, 8, 6, 4, 2))
		O = np.transpose(O, (0, 2, 4, 1, 3))
		O = np.reshape(O, (-1, 8*4))

		# O is (n_1*n_2*n_3 * C_B)
		O = np.matmul(O, np.reshape(A_4, (-1)))
		O = np.reshape(O, (-1, 2))

		# O is (n_1*n_2*n_3*n_4)
		O = np.matmul(O, B_4.T)
		O = np.reshape(O, (-1))
		
		return O

	X_1 = single_algo_2(X, A_1_1, B_1_1, A_1_2, B_1_2)
	X_2 = single_algo_2(X, A_2_1, B_2_1, A_2_2, B_2_2)
	X_3 = single_algo_2(X, A_3_1, B_3_1, A_3_2, B_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_2_2(Y, A_1_3, B_1_3, A_1_4, B_1_4)
	X_2 = single_algo_2_2(Y, A_2_3, B_2_3, A_2_4, B_2_4)
	X_3 = single_algo_2_2(Y, A_3_3, B_3_3, A_3_4, B_3_4)
	Y = X_1 + X_2 + X_3

	return Y


def divided_algo_1_reverse(XR):
	def single_algo_1(I, W_1, W_2):
		# W_1 to (m_1 * n_1*C)
		W_1 = np.reshape(np.reshape(W_1, (8, 6, -1)), (8, -1))

		# O is (m_4*m_3*m_2 * n_1*C)
		O = np.matmul(np.reshape(I, (8*8*8, 8)), W_1)

		# O is (n_1*m_4*m_3 * m_2*C)
		O = np.reshape(O, (-1, 8, 6, 8))
		O = np.transpose(O, (2, 0 , 1, 3))
		O = np.reshape(O, (-1, 8*8))

		# W_2 to (m_2*C * n_2)
		W_2 = np.reshape(W_2, (8, 6, -1))
		W_2 = np.transpose(W_2, (0, 2, 1))
		W_2 = np.reshape(W_2, (-1, 6))

		# O is (n_2*n_1*m_4 * m_3)
		O = np.reshape(np.matmul(O, W_2).T, (-1, 8))

		return O

	def single_algo_1_2(O, W_3, W_4):
		# W_3 to (m_3 * n_3*C)
		W_3 = np.reshape(np.reshape(W_3, (8, 6, -1)), (8, -1))

		# O is (n_2*n_1*m_4 * n_3*C)
		O = np.matmul(O, W_3)

		# O is (n_3*n_2*n_1 * m_4*C)
		O = np.reshape(O, (-1, 8, 6, 8))
		O = np.transpose(O, (2, 0 , 1, 3))
		O = np.reshape(O, (-1, 8*8))

		# W_4 to (m_4*C * n_4)
		W_4 = np.reshape(W_4, (8, 6, -1))
		W_4 = np.transpose(W_4, (0, 2, 1))
		W_4 = np.reshape(W_4, (-1, 6))

		# O is (n_4*n_3*n_2*n_1)
		O = np.reshape(np.matmul(O, W_4).T, (-1))

		return O

	AB_1_1 = np.kron(A_1_1, B_1_1)
	AB_1_2 = np.kron(A_1_2, B_1_2)
	AB_1_3 = np.kron(A_1_3, B_1_3)
	AB_1_4 = np.kron(A_1_4, B_1_4)
	AB_2_1 = np.kron(A_2_1, B_2_1)
	AB_2_2 = np.kron(A_2_2, B_2_2)
	AB_2_3 = np.kron(A_2_3, B_2_3)
	AB_2_4 = np.kron(A_2_4, B_2_4)
	AB_3_1 = np.kron(A_3_1, B_3_1)
	AB_3_2 = np.kron(A_3_2, B_3_2)
	AB_3_3 = np.kron(A_3_3, B_3_3)
	AB_3_4 = np.kron(A_3_4, B_3_4)

	X_1 = single_algo_1(XR, AB_1_1, AB_1_2)
	X_2 = single_algo_1(XR, AB_2_1, AB_2_2)
	X_3 = single_algo_1(XR, AB_3_1, AB_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_1_2(Y, AB_1_3, AB_1_4)
	X_2 = single_algo_1_2(Y, AB_2_3, AB_2_4)
	X_3 = single_algo_1_2(Y, AB_3_3, AB_3_4)
	Y = X_1 + X_2 + X_3

	return Y


def divided_algo_2_reverse(XR):
	def single_algo_2(I, A_1, B_1, A_2, B_2):
		# O is (m_4*m_3*m_2 * C_A)
		O = np.matmul(np.reshape(I, (8*8*8, 8)), A_1)

		# O is (n_1*m_4*m_3*m_2 * C_B*C_A)
		O = np.kron(B_1, O)

		# O is (n_1*m_4*m_3*C_B * m_2*C_A)
		O = np.reshape(O, (-1, 8, 2, 4))
		O = np.transpose(O, (0, 2, 1, 3))
		O = np.reshape(O, (-1, 8*4))

		# O is (n_1*m_4*m_3 * C_B)
		O = np.matmul(O, np.reshape(A_2, (-1)))
		O = np.reshape(O, (-1, 2))

		# O is (n_2*n_1*m_4 * m_3)
		O = np.matmul(O, B_2.T).T
		O = np.reshape(np.reshape(O, (-1)), (-1, 8))

		return O

	def single_algo_2_2(O, A_3, B_3, A_4, B_4):
		# O is (n_2*n_1*m_4 * C_A)
		O = np.matmul(O, A_3)

		# O is (n_3*n_2*n_1*m_4 * C_B*C_A)
		O = np.kron(B_3, O)

		# O is (n_3*n_2*n_1*C_B * m_4*C_A)
		O = np.reshape(O, (-1, 8, 2, 4))
		O = np.transpose(O, (0, 2, 1, 3))
		O = np.reshape(O, (-1, 8*4))

		# O is (n_3*n_2*n_1 * C_B)
		O = np.matmul(O, np.reshape(A_4, (-1)))
		O = np.reshape(O, (-1, 2))

		# O is (n_4*n_3*n_2*n_1)
		O = np.reshape(np.matmul(O, B_4.T).T, (-1))

		return O

	X_1 = single_algo_2(XR, A_1_1, B_1_1, A_1_2, B_1_2)
	X_2 = single_algo_2(XR, A_2_1, B_2_1, A_2_2, B_2_2)
	X_3 = single_algo_2(XR, A_3_1, B_3_1, A_3_2, B_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_2_2(Y, A_1_3, B_1_3, A_1_4, B_1_4)
	X_2 = single_algo_2_2(Y, A_2_3, B_2_3, A_2_4, B_2_4)
	X_3 = single_algo_2_2(Y, A_3_3, B_3_3, A_3_4, B_3_4)
	Y = X_1 + X_2 + X_3

	return Y


if __name__ == '__main__':
	s1 = time.time()
	Y_1 = normal_algo_1()
	e1 = time.time()

	s2 = time.time()
	Y_2 = divided_algo_1()
	e2 = time.time()

	s3 = time.time()
	Y_3 = divided_algo_2()
	e3 = time.time()

	XR = np.reshape(np.transpose(np.reshape(X, (8, 8, 8, 8)), (3, 2, 1, 0)), (-1))

	s4 = time.time()
	Y_4 = divided_algo_1_reverse(XR)
	e4 = time.time()
	Y_4 = np.reshape(np.transpose(np.reshape(Y_4, (6, 6, 6, 6)), (3, 2, 1, 0)), (-1))

	s5 = time.time()
	Y_5 = divided_algo_2_reverse(XR)
	e5 = time.time()
	Y_5 = np.reshape(np.transpose(np.reshape(Y_5, (6, 6, 6, 6)), (3, 2, 1, 0)), (-1))

	print('Normal (%lf) is:' % ((e1 - s1) * 1000))
	print(Y_1.shape)
	print(Y_1)
	print('\n')

	print('Divided 1 (%lf) is:' % ((e2 - s2) * 1000))
	print(Y_2.shape)
	print(Y_2)
	print('\n')

	print('Divided 2 (%lf) is:' % ((e3 - s3) * 1000))
	print(Y_3.shape)
	print(Y_3)
	print('\n')

	print('Divided 1 reversing (%lf) is:' % ((e4 - s4) * 1000))
	print(Y_4.shape)
	print(Y_4)
	print('\n')

	print('Divided 2 reversing (%lf) is:' % ((e5 - s5) * 1000))
	print(Y_5.shape)
	print(Y_5)
	print('\n')
